import cx_Oracle
import pandas as pd
from WindPy import w


def connect_db(dbname="wind") -> cx_Oracle.connect:
    db_config_dic = {
        'wind': "ofindataread/Ce_Sh1.Pa2s_ordb@29.157.144.5:19445/wdzxdb",
        'finchina': "findataread/Ce_Sh1.Pa2s_ordb@29.157.144.12:19412/fcdb",
        'gogal': "gogoalv2read/Ce_Sh1.Pa2s_ordb@29.157.144.5:19445/wdzxdb",
        'gil': "gildataread/gildataread$RFV518k@29.157.144.12:19412/fcdb",
        'LocalDB': "quant/123456@127.0.0.1:1521/XE"
    }
    # engine = create_engine('dialect+driver://username:password@host:port/database')
    db_config = db_config_dic.get(dbname, None)
    if db_config is not None:
        db_conn = cx_Oracle.connect(db_config)
        if db_conn:
            print("Database Access successfully!")
    else:
        db_conn = None
    return db_conn


class CBData:
    def __init__(self):
        # self._conn = connect_db()
        self.DB = dict({
            'Amt': pd.DataFrame(),
            'Close': pd.DataFrame(),
            'Outstanding': pd.DataFrame(),
        })

    def load_data(self, new_factor=None):
        data_csv = pd.read_csv('../data/CBondEODPrices.csv', header=0)
        data_csv['TRADE_DT'] = [str(i) for i in data_csv['TRADE_DT']]
        self.DB['Amt'] = data_csv.pivot(index='TRADE_DT', columns='S_INFO_WINDCODE', values='S_DQ_AMOUNT')
        self.DB['Close'] = data_csv.pivot(index='TRADE_DT', columns='S_INFO_WINDCODE', values='S_DQ_CLOSE')
        self.DB['Outstanding'] = data_csv.pivot(index='TRADE_DT', columns='S_INFO_WINDCODE', values='OUTSTANDING')
        if new_factor:
            self.DB[new_factor] = pd.DataFrame()
            codes = list(set(data_csv['S_INFO_WINDCODE']))
            start = min(data_csv['TRADE_DT'])
            end = max(data_csv['TRADE_DT'])
            w.start()
            if new_factor == 'Convprem':
                data_csv[new_factor] = w.wsd(codes, "convpremiumratio", start, end, usedf=True).Data
                self.DB[new_factor] = data_csv.pivot(index='TRADE_DT', columns='S_INFO_WINDCODE', values='Convprem')
            w.close()

    def exclude_special(self, hasEB=True):
        if hasEB:
            return list(self.DB['Outstanding'].columns)
        else:
            return None
